import fitz
from PIL import Image
import numpy as np
import os


def crop_white_borders(img):
    img_array = np.array(img)
    non_white = np.where(img_array != 255)
    if non_white[0].size == 0 or non_white[1].size == 0:
        return img
    left = np.min(non_white[1])
    right = np.max(non_white[1])
    return img.crop((left, 0, right, img.height))


def is_solid_color(img):
    img_array = np.array(img)
    if img_array.size == 0:
        return True
    first_pixel = img_array[0, 0]
    return np.all(img_array == first_pixel)


def has_black_pixel(img):
    pixels = img.getdata()
    for pixel in pixels:
        if pixel == (0, 0, 0):
            return True
    return False


def pdf_to_cropped_images(pdf_path, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    pdf_document = fitz.open(pdf_path)
    full_img = None
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)
        pix = page.get_pixmap()
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        img = crop_white_borders(img)
        if full_img is None:
            full_img = img
        else:
            new_img = Image.new('RGB', (full_img.width, full_img.height + img.height))
            new_img.paste(full_img, (0, 0))
            new_img.paste(img, (0, full_img.height))
            full_img = new_img

    width, height = full_img.size
    index = 0
    for y in range(0, height, 400):
        for x in range(0, width, 300):
            end_x = min(x + 300, width)
            end_y = min(y + 400, height)
            cropped_img = full_img.crop((x, y, end_x, end_y))
            cropped_width, _ = cropped_img.size
            if not is_solid_color(cropped_img) and (cropped_width >= 120 or has_black_pixel(cropped_img)):
                # 旋转图片 180 度
                cropped_img = cropped_img.rotate(180)
                output_path = os.path.join(output_folder, f"output_{index}.jpg")
                cropped_img.save(output_path, "JPEG")
                index += 1
    pdf_document.close()


if __name__ == "__main__":
    pdf_path = "test.pdf"
    output_folder = "test1"
    pdf_to_cropped_images(pdf_path, output_folder)
    
